"""
Diagnostic Logging Classes
==========================

Enhanced error logging with context and WebSocket diagnostics.
"""

import logging
import traceback
import uuid
from datetime import datetime
from typing import Dict, Any, Optional, List
from contextlib import contextmanager


class ContextualLogger:
    """
    Enhanced error logging with contextual information for debugging.
    """
    
    def __init__(self, name: str):
        self.logger = logging.getLogger(name)
        self.context: Dict[str, Any] = {}
        self.session_id = str(uuid.uuid4())[:8]
    
    def set_context(self, **kwargs) -> None:
        """Set persistent context for this logger instance."""
        self.context.update(kwargs)
    
    def clear_context(self) -> None:
        """Clear all context."""
        self.context.clear()
    
    @contextmanager
    def temp_context(self, **kwargs):
        """Temporarily add context for a block of code."""
        original_context = self.context.copy()
        self.context.update(kwargs)
        try:
            yield
        finally:
            self.context = original_context
    
    def debug(self, message: str, **extra_context) -> None:
        """Log debug message with context."""
        self._log_with_context(logging.DEBUG, message, **extra_context)
    
    def info(self, message: str, **extra_context) -> None:
        """Log info message with context."""
        self._log_with_context(logging.INFO, message, **extra_context)
    
    def warning(self, message: str, **extra_context) -> None:
        """Log warning message with context."""
        self._log_with_context(logging.WARNING, message, **extra_context)
    
    def error(self, message: str, exc_info=None, **extra_context) -> None:
        """Log error message with context and exception info."""
        self._log_with_context(logging.ERROR, message, exc_info=exc_info, **extra_context)
    
    def critical(self, message: str, exc_info=None, **extra_context) -> None:
        """Log critical message with context and exception info."""
        self._log_with_context(logging.CRITICAL, message, exc_info=exc_info, **extra_context)
    
    def error_with_context(self, message: str, exc_info=None, **extra_context) -> None:
        """Log error with full context and exception details."""
        context = {**self.context, **extra_context}
        
        error_data = {
            'context': context,
            'session_id': self.session_id,
            'timestamp': datetime.utcnow().isoformat()
        }
        
        if exc_info:
            error_data.update({
                'exception_type': type(exc_info).__name__ if hasattr(exc_info, '__name__') else str(type(exc_info)),
                'exception_message': str(exc_info),
                'stack_trace': traceback.format_exc()
            })
        
        self.logger.error(message, extra=error_data, exc_info=exc_info)
    
    def _log_with_context(self, level: int, message: str, exc_info=None, **extra_context) -> None:
        """Internal method to log with context."""
        context = {**self.context, **extra_context}
        
        log_data = {
            'context': context,
            'session_id': self.session_id,
            'timestamp': datetime.utcnow().isoformat()
        }
        
        self.logger.log(level, message, extra=log_data, exc_info=exc_info)


class WebSocketDiagnosticLogger:
    """
    Specialized logging for WebSocket connections and message flow.
    """
    
    def __init__(self, logger_name: str = 'capswriter.server.websocket.diagnostics'):
        self.logger = logging.getLogger(logger_name)
        self.connection_contexts: Dict[str, Dict[str, Any]] = {}
    
    def log_connection_event(self, 
                           event_type: str, 
                           client_id: str,
                           client_info: Optional[Dict[str, Any]] = None,
                           **extra_data) -> None:
        """Log WebSocket connection lifecycle events."""
        
        event_data = {
            'event_type': event_type,
            'client_id': client_id,
            'timestamp': datetime.utcnow().isoformat(),
            **extra_data
        }
        
        if client_info:
            event_data.update({
                'client_ip': client_info.get('remote_address', 'unknown'),
                'user_agent': client_info.get('user_agent', 'unknown'),
                'protocol': client_info.get('protocol', 'unknown')
            })
        
        # Store connection context for future reference
        if event_type == 'connection_established':
            self.connection_contexts[client_id] = {
                'established_at': datetime.utcnow().isoformat(),
                'client_info': client_info or {},
                'message_count': 0,
                'last_activity': datetime.utcnow().isoformat()
            }
        elif event_type == 'connection_closed':
            self.connection_contexts.pop(client_id, None)
        
        self.logger.info(f"websocket_{event_type}", extra=event_data)
    
    def log_message_flow(self, 
                        direction: str,
                        message_type: str, 
                        size: int,
                        client_id: str,
                        processing_time: Optional[float] = None,
                        content_preview: Optional[str] = None) -> None:
        """Log WebSocket message flow with detailed information."""
        
        # Update connection context
        if client_id in self.connection_contexts:
            context = self.connection_contexts[client_id]
            context['message_count'] += 1
            context['last_activity'] = datetime.utcnow().isoformat()
        
        message_data = {
            'direction': direction,
            'message_type': message_type,
            'size_bytes': size,
            'client_id': client_id,
            'timestamp': datetime.utcnow().isoformat()
        }
        
        if processing_time is not None:
            message_data['processing_time_ms'] = processing_time * 1000
        
        if content_preview:
            message_data['content_preview'] = content_preview[:100]  # Limit preview size
        
        self.logger.debug(f"websocket_message_{direction}", extra=message_data)
    
    def log_connection_error(self, 
                           client_id: str,
                           error_type: str,
                           error_message: str,
                           exc_info=None) -> None:
        """Log WebSocket connection errors with context."""
        
        error_data = {
            'client_id': client_id,
            'error_type': error_type,
            'error_message': error_message,
            'timestamp': datetime.utcnow().isoformat()
        }
        
        # Add connection context if available
        if client_id in self.connection_contexts:
            error_data['connection_context'] = self.connection_contexts[client_id]
        
        if exc_info:
            error_data['stack_trace'] = traceback.format_exc()
        
        self.logger.error("websocket_connection_error", extra=error_data, exc_info=exc_info)
    
    def log_protocol_error(self,
                          client_id: str,
                          expected: str,
                          received: str,
                          raw_data: Optional[bytes] = None) -> None:
        """Log WebSocket protocol errors."""
        
        protocol_data = {
            'client_id': client_id,
            'expected': expected,
            'received': received,
            'timestamp': datetime.utcnow().isoformat()
        }
        
        if raw_data:
            # Log hex representation of raw data for debugging
            protocol_data['raw_data_hex'] = raw_data[:50].hex()  # Limit to first 50 bytes
            protocol_data['raw_data_size'] = len(raw_data)
        
        self.logger.error("websocket_protocol_error", extra=protocol_data)
    
    def get_connection_summary(self) -> Dict[str, Any]:
        """Get summary of current WebSocket connections."""
        
        summary = {
            'active_connections': len(self.connection_contexts),
            'connections': list(self.connection_contexts.keys()),
            'timestamp': datetime.utcnow().isoformat()
        }
        
        # Add per-connection summaries
        connection_summaries = []
        for client_id, context in self.connection_contexts.items():
            connection_summaries.append({
                'client_id': client_id,
                'established_at': context['established_at'],
                'message_count': context['message_count'],
                'last_activity': context['last_activity']
            })
        
        summary['connection_details'] = connection_summaries
        
        self.logger.info("websocket_connection_summary", extra=summary)
        return summary


class AudioDiagnosticLogger:
    """
    Specialized logging for audio processing diagnostics.
    """
    
    def __init__(self, logger_name: str = 'capswriter.audio.diagnostics'):
        self.logger = logging.getLogger(logger_name)
    
    def log_audio_stream_event(self,
                              event_type: str,
                              stream_info: Dict[str, Any],
                              **extra_data) -> None:
        """Log audio stream lifecycle events."""
        
        event_data = {
            'event_type': event_type,
            'stream_info': stream_info,
            'timestamp': datetime.utcnow().isoformat(),
            **extra_data
        }
        
        self.logger.info(f"audio_stream_{event_type}", extra=event_data)
    
    def log_audio_processing_error(self,
                                 error_type: str,
                                 audio_info: Dict[str, Any],
                                 error_message: str,
                                 exc_info=None) -> None:
        """Log audio processing errors."""
        
        error_data = {
            'error_type': error_type,
            'audio_info': audio_info,
            'error_message': error_message,
            'timestamp': datetime.utcnow().isoformat()
        }
        
        if exc_info:
            error_data['stack_trace'] = traceback.format_exc()
        
        self.logger.error("audio_processing_error", extra=error_data, exc_info=exc_info)
    
    def log_recognition_debug(self,
                            audio_segment_info: Dict[str, Any],
                            model_output: str,
                            confidence_scores: Optional[List[float]] = None) -> None:
        """Log detailed recognition debugging information."""
        
        debug_data = {
            'audio_segment_info': audio_segment_info,
            'model_output': model_output,
            'output_length': len(model_output),
            'timestamp': datetime.utcnow().isoformat()
        }
        
        if confidence_scores:
            debug_data.update({
                'confidence_scores': confidence_scores,
                'avg_confidence': sum(confidence_scores) / len(confidence_scores),
                'min_confidence': min(confidence_scores),
                'max_confidence': max(confidence_scores)
            })
        
        self.logger.debug("recognition_debug", extra=debug_data)


class ModelDiagnosticLogger:
    """
    Specialized logging for model loading and inference diagnostics.
    """
    
    def __init__(self, logger_name: str = 'capswriter.model.diagnostics'):
        self.logger = logging.getLogger(logger_name)
    
    def log_model_loading(self,
                         model_name: str,
                         model_path: str,
                         loading_time: float,
                         model_size: Optional[int] = None) -> None:
        """Log model loading information."""
        
        loading_data = {
            'model_name': model_name,
            'model_path': model_path,
            'loading_time_seconds': loading_time,
            'timestamp': datetime.utcnow().isoformat()
        }
        
        if model_size:
            loading_data['model_size_mb'] = model_size / (1024 * 1024)
        
        self.logger.info("model_loaded", extra=loading_data)
    
    def log_model_error(self,
                       model_name: str,
                       error_type: str,
                       error_message: str,
                       exc_info=None) -> None:
        """Log model-related errors."""
        
        error_data = {
            'model_name': model_name,
            'error_type': error_type,
            'error_message': error_message,
            'timestamp': datetime.utcnow().isoformat()
        }
        
        if exc_info:
            error_data['stack_trace'] = traceback.format_exc()
        
        self.logger.error("model_error", extra=error_data, exc_info=exc_info)
    
    def log_inference_metrics(self,
                            model_name: str,
                            input_shape: tuple,
                            inference_time: float,
                            memory_usage: Optional[float] = None) -> None:
        """Log model inference performance metrics."""
        
        inference_data = {
            'model_name': model_name,
            'input_shape': input_shape,
            'inference_time_ms': inference_time * 1000,
            'timestamp': datetime.utcnow().isoformat()
        }
        
        if memory_usage:
            inference_data['memory_usage_mb'] = memory_usage
        
        self.logger.debug("model_inference", extra=inference_data)